"""
Helper functions for sql database

>>> # get the test database
>>> uri = get_testuri()
>>> engine = get_engine( uri )
>>>
>>> # testing the SQL Expression Language 
>>> id = make_label( engine, name='TEST-LABEL', access='R', clobber=True)
>>> type(id) == int
True
"""
import csv
import atlas
from atlas import util
from sqlalchemy import *  
from sqlalchemy.sql import func, select
from sqlalchemy.exceptions import *  
from sqlalchemy.orm import mapper, sessionmaker, relation, scoped_session
from sqlalchemy import pool
from urlparse import urlparse

metadata = MetaData()

label_table = Table('labels', metadata,
    Column('id', Integer, primary_key=True),
    Column('name', String(50),  unique=True),  # identify the label
    Column('desc', String(100)), # help on the label
    Column('access', String(1)), # read/write
    CheckConstraint("access in( 'R', 'W')", name='check1')
)

class Label(object):
    def __init__(self, name, desc='', access='W'):
        self.name = name
        self.desc = desc
        self.access = access
    def __repr__(self):
        return 'Label: name=%s, access=%s' %(self.name, self.access)

feature_table = Table('feature', metadata,
    Column('id', Integer, primary_key=True),
    Column('name', String(50), index=True),
    Column('altname', String(50), index=True),
    Column('chrom', String(50), index=True),
    Column('strand', String(1), index=True),
    Column('start', Integer,index=True ),
    Column('end', Integer, index=True ),
    Column('freetext', String,),
    Column('value', Float),
    Column('label_id', Integer, ForeignKey('labels.id')),
)

class Feature(object):
    def __init__(self, name, chrom, start, end, json='[]', value=0):
        self.name, self.chrom, self.start, self.end = name, chrom, start, end
        self.json, self.value = json, value
    def __repr__(self):
        return 'Feature: name=%s, label=%s' %( self.name, self.label.name )
 
mapper(Feature, feature_table) 
mapper(Label, label_table, properties={'features':relation(Feature, backref='label', cascade="all, delete, delete-orphan") } ) 

def get_session( uri ):
    """
    Returns an sql session        
    """
    engine = get_engine( uri )
    session = scoped_session( sessionmaker( bind=engine, autoflush=True ) )
    return session

def get_testuri( fname='test-database.sqlite' ):
    """
    Returns a universal resource identifier for a test database (sqlite)
    """
    uri = atlas.testdata( fname, verify=False )
    return normalize_uri( uri )

def get_engine( uri, create=True ):
    """
    Returns an engine based on a universal resource identifier. 
    """
    global metadata
   
    atlas.debug( 'uri=%s' % uri )
    engine = create_engine( uri )
    if create:
        metadata.create_all(engine) 
    return engine

def select_features( engine, label_id):
    """
    Creates a new label. Drops existing data with this label
    if the label existed before.
    """
    table = feature_table
    conn  = engine.connect()
    sel   = select( [ table ], table.c.label_id==label_id )
    rows  = conn.execute( sel )
    return rows

def bulk_insert( engine, values, table=feature_table ):
    "Bulk inserts into the table"
    atlas.debug( 'inserting %s rows' % atlas.commify(len(values)) )
    if not values:
        return
    insert = table.insert()
    conn = engine.connect()
    conn.execute( insert, values)
    conn.close()

def bulk_delete( engine, label_id, table=feature_table ):
    "Bulk delete features"
    conn = engine.connect()
    lt = label_table
    count = conn.execute( table.count( table.c.label_id == label_id )).scalar()
    atlas.info( 'deleting %s rows' % atlas.commify(count) )
    conn.execute( table.delete( table.c.label_id == label_id ) )
    conn.execute( lt.delete( lt.c.id == label_id ) )
    conn.close()

def make_label( engine, name, access='R', clobber=False):
    """
    Creates a new label. Drops existing data with this label
    if the label existed before.
    """
    labt  = label_table
    conn  = engine.connect()
    sel   = select( [ labt  ], labt.c.name==name )
    rows  = conn.execute( sel )
    first = rows.fetchone()
    if first:
        if not clobber:
            raise Exception('label %s exists, set clobber to True to override' % name )
        atlas.info( "label %s exists, bulk deleting" % name )
        bulk_delete(engine=engine, label_id=first.id)

    ins  = labt.insert(values=dict(name=name, access=access, text='') )
    rows = conn.execute( ins )

    # reselect the label
    rows  = conn.execute( sel )
    first = rows.fetchone()
    return first[0]

@util.stop_on_error
def load_features( sql_uri, file_name, clobber=False, drop=False ):
    """
    Loads features into the database
    """
    atlas.info('loading features from %s' % file_name )

    # load up all features into memory
    reader = util.dict_reader(  file_name )
    lineno, first  = reader.next()
    label_name = first['label']
    
    engine = get_engine( sql_uri )
        
    # inserts label for the peak, if clobber is True will remove all data for this label!
    try:
        label_id = make_label( engine, name=label_name, clobber=clobber )
    except Exception, exc:
        atlas.warn('feature insert aborted %s' % exc )
        return

    # drop sql indices (speeds up bulk inserting) 
    if drop:
        drop_indices(engine)

    reader = util.dict_reader(  file_name )
    values = []
    for lineno, row in reader:
        row['label_id'] = label_id
        values.append( row )
    
    bulk_insert( engine=engine, values=values )
    
    if drop:
        create_indices(engine)
      
def drop_indices( engine, table=feature_table):
    "Shortcut to dropping indices"
    atlas.info('dropping feature indices')
    manage_indices( engine, table=table, cmd='drop')

def create_indices( engine, table=feature_table):
    "Shortcut to creating indices"
    atlas.info('creating feature indices')
    manage_indices( engine, table=table, cmd='create')

def manage_indices( engine, table, cmd='drop'):
    "Drops all the indices"
    for index in table.indexes:
        try:
            if cmd == 'drop':
                index.drop( engine )
            else:
                index.create( engine )
        except OperationalError, exc:
            # already done
            pass

def normalize_uri(uri):
    """Normalizes a windows filepath into a sqlite uri"""
    
    uri = uri.replace("\\", "/") # another windows fix
    elems = uri.split(':')
    if len(elems)==1:
        # no scheme passed, assume SQLite and build the proper uri
        uri = "sqlite:///" + uri
    else:
        if elems[0].upper() in "CDEFG": # windows drives
            uri = "sqlite:///" + uri[2:]
   
    return uri
    
def test(verbose=0):
    """
    Main testrunnner
    """
    import doctest
    doctest.testmod( optionflags=doctest.ELLIPSIS + doctest.NORMALIZE_WHITESPACE, verbose=verbose )

if __name__ == '__main__':
    test(verbose=0)
    if 0:
        uri = get_testuri()
        engine = get_engine( uri )
        
        make_label( engine=engine, name='TEST-LABEL', access='R', clobber=True)
